from keras.layers import Dropout,add,Input,Conv2D,BatchNormalization,Activation,Reshape,Dense,MaxPool2D
from keras import Model
import tensorflow as tf
from keras.utils.vis_utils import plot_model


def res_block(input):
    # assert
    img_input = input
    # 由于网络层次较浅，容易欠拟合，正则化视情况而定
    img_input = BatchNormalization()(img_input)  
    act1= Activation('relu')(img_input)
    conv1 = Conv2D(filters=act1.shape[-1].value,
                    kernel_size=3,
                    strides=1,
                    padding='same',
                    kernel_initializer='glorot_normal',
                    bias_initializer='zeros'
                    )(act1)
    block_out = add([conv1,img_input])
    block_out_act = Activation('relu')(block_out)
    return block_out_act

class Nets_2D:
    '''
    初始化类的时候🙏输入张量[None,width,height,channels]
    或者keras的Input层
    '''
    def __init__(self,width,height):
        self.input = Input(shape=(width,height,1))

    def mynet_1(self):
        # 其实就是LeNet-5，似乎是手写体专用网络
        # 输入为Input(shape=(width,height,1))
        real_input = self.input
        # 归一化
        bn1 = BatchNormalization()(real_input)
        cv1 = Conv2D(filters=6, kernel_size=3, strides=1, activation='relu', padding='same')(bn1)
        pl1 = MaxPool2D(strides=(2,2))(cv1)
        cv2 = Conv2D(filters=16, kernel_size=3, strides=1, activation='relu', padding='same')(pl1)
        pl2 = MaxPool2D(strides=(2,2))(cv2)
        tem = Reshape((pl2.shape[1].value*pl2.shape[2].value*pl2.shape[3].value,))(pl2)
        fu1 = Dense(120, activation='relu')(tem)
        fu2 = Dense(84, activation='relu')(fu1)
        dp1 = Dropout(0.85)(fu2)
        out = Dense(10, activation='softmax')(dp1)

        model = Model(inputs=real_input,outputs=out)
        return model

    def mynet_2(self):
        # 超级全连接层，🏠🛣搞的残差块，
        real_input = self.input
        # 层数浅就不做归一化了
        res1 = res_block(real_input)
        cv1  = Conv2D(filters=6, kernel_size=3, strides=2, activation='relu', padding='same')(res1)
        res2 = res_block(cv1)
        cv2  = Conv2D(filters=16,kernel_size=3, strides=2, activation='relu', padding='same')(res2)
        tem  = Reshape((real_input.shape[1].value*real_input.shape[2].value*real_input.shape[3].value,))(cv2)
        ful1 = Dense(392, activation='relu')(tem)
        ful2 = Dense(196, activation='relu')(ful1)
        ful3 = Dense(98, activation='relu')(ful2)
        dp1  = Dropout(0.85)(ful3)
        out  = Dense(10, activation='softmax')(dp1)

        model = Model(inputs=real_input,outputs=out)
        return model

    # def vgg_16_(self):


if __name__ == "__main__":
    input_height = 28
    input_width = 28
    Nets = Nets_2D(input_width,input_height)
    # 自定义的LetNet-5，真的是个铁🦐🦐
    plot_model(Nets.mynet_1(),to_file='mynet_1.png',show_shapes=True)
    plot_model(Nets.mynet_2(),to_file='mynet_2.png',show_shapes=True)